"""
Author: xubing
Date: 2024-11-21
Description: 百川大模型客户端
"""

from http import HTTPStatus
from typing import Generator, List

import dashscope
from dashscope.api_entities.dashscope_response import DashScopeAPIResponse
from decorators import timeit
from logger_config_db import setup_logger

from myUtils.clients.base_client import BaseClient, ChatResponse, Message

# 设置日志记录器
logger = setup_logger(__name__)


class BaiChuanClient(BaseClient):
    """百川大模型客户端"""

    def __init__(self, model_name):
        super().__init__(model_name)

    @timeit
    def chat(self, messages) -> ChatResponse:
        try:
            dashscope.api_key = "sk-2eea77e26a464c25b91d2e6a45aebd55"
            response = dashscope.Generation.call(
                model=self.model_name,
                messages=messages,
                result_format="message",
                stream=False,
            )
            if response.status_code == HTTPStatus.OK:
                return ChatResponse(
                    content=response.output.choices[0].message["content"])
            else:
                error_msg = (
                    "Request id: %s, Status code: %s, error code: %s, error message: %s"
                    % (
                        response.request_id,
                        response.status_code,
                        response.code,
                        response.message,
                    ))
                return ChatResponse(content="",
                                    status=False,
                                    error_msg=error_msg)
        except Exception as e:
            logger.error(f"Chat error with {self.model_name}: {e}")
            return ChatResponse(content="", status=False, error_msg=str(e))

    def chat_stream(self, messages) -> Generator[str, None, None]:
        try:
            dashscope.api_key = "sk-2eea77e26a464c25b91d2e6a45aebd55"
            response = dashscope.Generation.call(
                model=self.model_name,
                messages=messages,
                result_format="message",
                stream=True,
            )

            if isinstance(response, DashScopeAPIResponse):
                # 单次响应，非流式
                if response.status_code == HTTPStatus.OK:
                    yield response.output.choices[0].message["content"]
                else:
                    yield f"Error: {response.message}"
            else:
                # 流式响应
                for chunk in response:
                    if chunk.status_code == HTTPStatus.OK and chunk.output and chunk.output.choices:
                        if "content" in chunk.output.choices[0].message:
                            yield chunk.output.choices[0].message["content"]
                    elif chunk.status_code != HTTPStatus.OK:
                        yield f"Error: {chunk.message}"
        except Exception as e:
            logger.error(f"Stream chat error with {self.model_name}: {e}")
            yield f"Error: {str(e)}"
